<template>
  <div class="model-train-container">
    <el-steps :active="step" finish-status="success" align-center>
      <el-step title="项目创建"></el-step>
      <el-step title="网络架构设计"></el-step>
      <el-step title="参数配置"></el-step>
      <el-step title="模型训练"></el-step>
      <el-step title="模型评估"></el-step>
    </el-steps>

    <!-- Step 1: 项目创建 -->
    <div v-show="step === 0" class="step-content">
      <el-form :model="project" label-width="100px" class="form">
        <el-form-item label="模型名称">
          <el-input v-model="project.name" placeholder="请输入模型名称"></el-input>
        </el-form-item>
        <el-form-item label="深度学习框架">
          <el-select v-model="project.framework" placeholder="请选择">
            <el-option label="TensorFlow" value="TensorFlow"></el-option>
            <el-option label="PyTorch" value="PyTorch"></el-option>
          </el-select>
        </el-form-item>
        <el-form-item label="上传数据">
          <el-upload
            class="upload-demo"
            drag
            :auto-upload="false"
            :on-change="handleFileChange"
            :file-list="fileList"
            accept=".csv,.json"
          >
            <i class="el-icon-upload"></i>
            <div class="el-upload__text">将文件拖到此处，或<em>点击上传</em></div>
            <div class="el-upload__tip" slot="tip">支持 .csv、.json 格式</div>
          </el-upload>
        </el-form-item>
        <el-form-item label="训练/验证集比例">
          <el-slider v-model="project.split" :min="6" :max="9" :step="1" show-stops>
          </el-slider>
          <div>训练集: {{project.split}}0% 验证集: {{(10-project.split)}}0%</div>
        </el-form-item>
      </el-form>
      <div class="btns">
        <el-button type="primary" @click="nextStep" :disabled="!canNextStep1">下一步</el-button>
      </div>
    </div>

    <!-- Step 2: 网络架构设计 -->
    <div v-show="step === 1" class="step-content">
      <div class="network-design">
        <div class="toolbox">
          <div class="tool-title">组件工具箱</div>
          <el-card class="tool-item" draggable @dragstart.native="onDragStart('LSTM')">
            <div><i class="el-icon-cpu"></i> LSTM层</div>
            <div class="tool-desc">循环神经网络</div>
          </el-card>
          <el-card class="tool-item" draggable @dragstart.native="onDragStart('Dropout')">
            <div><i class="el-icon-refresh"></i> Dropout层</div>
            <div class="tool-desc">防止过拟合</div>
          </el-card>
        </div>
        <div
          class="design-area"
          @dragover.prevent
          @drop="onDrop"
        >
          <div class="design-title">网络架构设计</div>
          <div v-if="modelLayers.length === 0" class="design-placeholder">
            拖拽组件到此处进行模型设计
          </div>
          <div v-else class="layer-list">
            <div
              v-for="(layer, idx) in modelLayers"
              :key="idx"
              class="layer-item"
            >
              <span>{{layer.type}}</span>
              <template v-if="layer.type==='LSTM'">
                <el-input-number v-model="layer.units" :min="1" :max="512" size="mini" label="神经元"></el-input-number>
                <span>神经元</span>
              </template>
              <template v-if="layer.type==='Dropout'">
                <el-input-number v-model="layer.rate" :min="0.1" :max="0.9" :step="0.05" size="mini" label="比率" :precision="2"></el-input-number>
                <span>比率</span>
              </template>
              <el-button type="text" icon="el-icon-delete" @click="removeLayer(idx)"></el-button>
            </div>
          </div>
        </div>
      </div>
      <div class="btns">
        <el-button @click="prevStep">上一步</el-button>
        <el-button type="primary" @click="nextStep" :disabled="!canNextStep2">下一步</el-button>
      </div>
    </div>

    <!-- Step 3: 参数配置 -->
    <div v-show="step === 2" class="step-content">
      <el-form :model="params" label-width="100px" class="form">
        <el-form-item label="学习率">
          <el-input-number v-model="params.learningRate" :min="0.0001" :max="0.1" :step="0.0001" :precision="4"></el-input-number>
        </el-form-item>
        <el-form-item label="批量大小">
          <el-input-number v-model="params.batchSize" :min="8" :max="256" :step="8"></el-input-number>
        </el-form-item>
        <el-form-item label="迭代次数">
          <el-input-number v-model="params.epochs" :min="10" :max="1000" :step="10"></el-input-number>
        </el-form-item>
        <el-form-item label="早停策略">
          <el-switch v-model="params.earlyStop"></el-switch>
          <span style="margin-left:10px;">验证集损失连续5轮不下降则停止</span>
        </el-form-item>
      </el-form>
      <div class="btns">
        <el-button @click="prevStep">上一步</el-button>
        <el-button type="primary" @click="nextStep">下一步</el-button>
      </div>
    </div>

    <!-- Step 4: 模型训练 -->
    <div v-show="step === 3" class="step-content">
      <div class="train-info">
        <div>模型名称：{{project.name}}</div>
        <div>框架：{{project.framework}}</div>
        <div>网络结构：
          <span v-for="(layer, idx) in modelLayers" :key="idx">
            {{layer.type}}<template v-if="layer.type==='LSTM'">({{layer.units}})</template>
            <template v-if="layer.type==='Dropout'">({{layer.rate}})</template>
            <span v-if="idx < modelLayers.length-1">→</span>
          </span>
        </div>
        <div>参数：学习率{{params.learningRate}}，批量{{params.batchSize}}，迭代{{params.epochs}}</div>
      </div>
      <div class="train-charts">
        <div id="lossChart" style="width: 45%; height: 250px; display:inline-block;"></div>
        <div id="accChart" style="width: 45%; height: 250px; display:inline-block;"></div>
      </div>
      <div class="btns">
        <el-button @click="prevStep" :disabled="training">上一步</el-button>
        <el-button type="primary" @click="startTrain" :disabled="training || trained">开始训练</el-button>
        <el-button v-if="trained" type="success" @click="nextStep">训练完成，下一步</el-button>
      </div>
    </div>

    <!-- Step 5: 模型评估 -->
    <div v-show="step === 4" class="step-content">
      <div class="eval-section">
        <div id="cmChart" style="width: 45%; height: 250px; display:inline-block;"></div>
        <div class="eval-metrics">
          <div>准确率：<span class="metric-value">{{(metrics.accuracy*100).toFixed(1)}}%</span></div>
          <div>精确率：<span class="metric-value">{{(metrics.precision*100).toFixed(1)}}%</span></div>
          <div>召回率：<span class="metric-value">{{(metrics.recall*100).toFixed(1)}}%</span></div>
        </div>
      </div>
      <div class="btns">
        <el-button @click="prevStep">上一步</el-button>
        <el-button type="primary" @click="downloadWeights">导出模型（.h5）</el-button>
      </div>
    </div>
  </div>
</template>

<script>
import { Message } from 'element-ui'
import * as echarts from 'echarts'

export default {
  name: 'ModelTraining',
  data() {
    return {
      step: 0,
      project: {
        name: '',
        framework: '',
        split: 8, // 8:2
        data: null,
        raw: null
      },
      fileList: [],
      modelLayers: [],
      dragLayer: null,
      params: {
        learningRate: 0.001,
        batchSize: 64,
        epochs: 200,
        earlyStop: true
      },
      training: false,
      trained: false,
      trainHistory: {
        loss: [],
        acc: []
      },
      metrics: {
        accuracy: 0.952,
        precision: 0.948,
        recall: 0.955
      },
      confusionMatrix: [
        [95, 5],
        [3, 97]
      ]
    }
  },
  computed: {
    canNextStep1() {
      return this.project.name && this.project.framework && this.project.data
    },
    canNextStep2() {
      // 必须有LSTM和Dropout层
      return (
        this.modelLayers.length >= 2 &&
        this.modelLayers.some(l => l.type === 'LSTM') &&
        this.modelLayers.some(l => l.type === 'Dropout')
      )
    }
  },
  mounted() {
    // 恢复localStorage
    const saved = localStorage.getItem('model-train-data')
    if (saved) {
      const d = JSON.parse(saved)
      Object.assign(this.$data, d)
    }
    if (this.step === 3 && this.trainHistory.loss.length) {
      this.$nextTick(this.renderTrainCharts)
    }
    if (this.step === 4) {
      this.$nextTick(this.renderCMChart)
    }
  },
  watch: {
    step(val) {
      if (val === 3 && this.trainHistory.loss.length) {
        this.$nextTick(this.renderTrainCharts)
      }
      if (val === 4) {
        this.$nextTick(this.renderCMChart)
      }
    }
  },
  methods: {
    nextStep() {
      this.saveToLocal()
      this.step++
      if (this.step === 3 && this.trainHistory.loss.length) {
        this.$nextTick(this.renderTrainCharts)
      }
      if (this.step === 4) {
        this.$nextTick(this.renderCMChart)
      }
    },
    prevStep() {
      this.step--
      if (this.step === 3 && this.trainHistory.loss.length) {
        this.$nextTick(this.renderTrainCharts)
      }
      if (this.step === 4) {
        this.$nextTick(this.renderCMChart)
      }
    },
    handleFileChange(file) {
      this.fileList = [file]
      const reader = new FileReader()
      reader.onload = e => {
        this.project.data = e.target.result
        this.project.raw = file.name
        Message.success('数据上传成功（mock数据，实际未解析）')
      }
      reader.readAsText(file.raw)
    },
    onDragStart(type) {
      this.dragLayer = type
    },
    onDrop() {
      if (this.dragLayer === 'LSTM') {
        this.modelLayers.push({ type: 'LSTM', units: 128 })
      } else if (this.dragLayer === 'Dropout') {
        this.modelLayers.push({ type: 'Dropout', rate: 0.2 })
      }
      this.dragLayer = null
    },
    removeLayer(idx) {
      this.modelLayers.splice(idx, 1)
    },
    saveToLocal() {
      localStorage.setItem(
        'model-train-data',
        JSON.stringify({
          step: this.step,
          project: this.project,
          fileList: this.fileList,
          modelLayers: this.modelLayers,
          params: this.params,
          training: this.training,
          trained: this.trained,
          trainHistory: this.trainHistory,
          metrics: this.metrics,
          confusionMatrix: this.confusionMatrix
        })
      )
    },
    startTrain() {
      if (this.training) return
      this.training = true
      this.trained = false
      this.trainHistory = { loss: [], acc: [] }
      // mock训练过程
      let loss = 0.8
      let acc = 0.5
      let bestValLoss = 0.8
      let noImprove = 0
      let valLoss = 0.8
      const epochs = this.params.epochs
      const history = { loss: [], acc: [] }
      let i = 0
      const trainStep = () => {
        if (i >= epochs) {
          this.finishTrain(history)
          return
        }
        // mock loss/acc
        loss = Math.max(0.1, loss - Math.random() * 0.1)
        acc = Math.min(1, acc + Math.random() * 0.05)
        // mock val loss
        valLoss = loss + (Math.random() - 0.5) * 0.05
        history.loss.push(Number(loss.toFixed(4)))
        history.acc.push(Number(acc.toFixed(4)))
        // 早停
        if (this.params.earlyStop) {
          if (valLoss < bestValLoss) {
            bestValLoss = valLoss
            noImprove = 0
          } else {
            noImprove++
          }
          if (noImprove >= 5) {
            this.finishTrain(history)
            return
          }
        }
        i++
        this.trainHistory = { ...history }
        this.$nextTick(this.renderTrainCharts)
        setTimeout(trainStep, 80)
      }
      trainStep()
    },
    finishTrain(history) {
      this.training = false
      this.trained = true
      this.trainHistory = { ...history }
      // mock评估
      this.metrics = {
        accuracy: 0.952 + Math.random() * 0.02,
        precision: 0.948 + Math.random() * 0.02,
        recall: 0.955 + Math.random() * 0.02
      }
      this.confusionMatrix = [
        [95 + Math.floor(Math.random() * 3), 5 - Math.floor(Math.random() * 2)],
        [3 + Math.floor(Math.random() * 2), 97 + Math.floor(Math.random() * 2)]
      ]
      this.saveToLocal()
      Message.success('训练完成！')
    },
    renderTrainCharts() {
      // 损失曲线
      const lossChart = echarts.init(document.getElementById('lossChart'))
      lossChart.setOption({
        title: { text: '损失曲线', left: 'center', top: 10, textStyle: { fontSize: 14 } },
        xAxis: { type: 'category', data: this.trainHistory.loss.map((_, i) => i + 1) },
        yAxis: { type: 'value', min: 0, max: 1 },
        series: [
          {
            name: 'loss',
            type: 'line',
            data: this.trainHistory.loss,
            smooth: true,
            lineStyle: { color: '#409EFF' }
          }
        ]
      })
      // 准确率曲线
      const accChart = echarts.init(document.getElementById('accChart'))
      accChart.setOption({
        title: { text: '准确率', left: 'center', top: 10, textStyle: { fontSize: 14 } },
        xAxis: { type: 'category', data: this.trainHistory.acc.map((_, i) => i + 1) },
        yAxis: { type: 'value', min: 0, max: 1 },
        series: [
          {
            name: 'accuracy',
            type: 'line',
            data: this.trainHistory.acc,
            smooth: true,
            lineStyle: { color: '#67C23A' }
          }
        ]
      })
    },
    renderCMChart() {
      // 混淆矩阵
      const cmChart = echarts.init(document.getElementById('cmChart'))
      cmChart.setOption({
        title: { text: '混淆矩阵', left: 'center', top: 10, textStyle: { fontSize: 14 } },
        tooltip: { trigger: 'item', formatter: params => `预测${params.data[1]} 实际${params.data[0]}: ${params.data[2]}` },
        xAxis: { type: 'category', data: ['正常', '故障'] },
        yAxis: { type: 'category', data: ['正常', '故障'] },
        visualMap: {
          min: 0,
          max: 100,
          calculable: true,
          orient: 'horizontal',
          left: 'center',
          bottom: '5%'
        },
        series: [
          {
            name: '混淆矩阵',
            type: 'heatmap',
            data: [
              [0, 0, this.confusionMatrix[0][0]],
              [0, 1, this.confusionMatrix[0][1]],
              [1, 0, this.confusionMatrix[1][0]],
              [1, 1, this.confusionMatrix[1][1]]
            ],
            label: { show: true }
          }
        ]
      })
    },
    downloadWeights() {
      // mock权重内容
      const content = 'mock h5 weights'
      const blob = new Blob([content], { type: 'application/octet-stream' })
      const a = document.createElement('a')
      a.href = URL.createObjectURL(blob)
      a.download = `${this.project.name || 'model'}.h5`
      a.click()
      Message.success('模型权重已导出（mock文件）')
    }
  }
}
</script>

<style scoped>
.model-train-container {
  background: #fff;
  padding: 32px 32px 0 32px;
  min-height: 100vh;
}
.step-content {
  margin: 32px auto 0 auto;
  max-width: 900px;
  background: #fafbfc;
  border-radius: 8px;
  padding: 32px 40px 24px 40px;
  box-shadow: 0 2px 8px #f0f1f2;
}
.form {
  max-width: 600px;
}
.btns {
  margin-top: 32px;
  text-align: right;
}
.network-design {
  display: flex;
  flex-direction: row;
  justify-content: space-between;
}
.toolbox {
  width: 180px;
  margin-right: 32px;
}
.tool-title {
  font-weight: bold;
  margin-bottom: 12px;
}
.tool-item {
  margin-bottom: 16px;
  cursor: grab;
}
.tool-desc {
  font-size: 12px;
  color: #888;
}
.design-area {
  flex: 1;
  min-height: 220px;
  border: 2px dashed #d3d3d3;
  border-radius: 8px;
  background: #f6f8fa;
  padding: 16px;
  position: relative;
}
.design-title {
  font-weight: bold;
  margin-bottom: 8px;
}
.design-placeholder {
  color: #bbb;
  text-align: center;
  margin-top: 40px;
}
.layer-list {
  display: flex;
  flex-direction: column;
  gap: 12px;
}
.layer-item {
  background: #fff;
  border-radius: 4px;
  padding: 8px 12px;
  display: flex;
  align-items: center;
  gap: 8px;
  box-shadow: 0 1px 4px #eee;
}
.train-info {
  margin-bottom: 24px;
  font-size: 15px;
  color: #333;
}
.train-charts {
  display: flex;
  flex-direction: row;
  justify-content: space-between;
  gap: 24px;
}
.eval-section {
  display: flex;
  flex-direction: row;
  justify-content: space-between;
  align-items: flex-start;
}
.eval-metrics {
  margin-left: 32px;
  font-size: 18px;
  color: #333;
}
.metric-value {
  color: #67C23A;
  font-weight: bold;
  font-size: 22px;
  margin-left: 8px;
}
</style>
